#include <stdio.h>
#include <stdint.h>
#include <fstream>
#include <vector>
#include <sstream>

void dumpbin(void *buf, size_t sz, const std::string &path) {
  std::ofstream ofs;
  ofs.open(path, std::ios::out | std::ios::trunc | std::ios::binary);
  ofs.write((const char *)buf, sz);
  // printf("# dump %s \n", path.c_str());
}

template<typename T>
void memset(void *buf, const T val, int sz) {
  for(int id = 0; id < sz; ++id){
    ((T*)buf)[id] = val;
  }
}

void mp_gemm_w1a8_host(
  int m, int n, int k, void *A, void *B, void *C) {
  const int bk = 8;

  for(int idm=0; idm < m; ++idm){
    for(int idn=0; idn < n; ++idn){
      int32_t dot = 0;
      int32_t *C_int32 = ((int32_t*)C)+idn*m+idm;
      for(int idbk = 0; idbk < k/bk; ++idbk){
        uint8_t A_uint1 = ((uint8_t*)A)[idm*(k/bk)+idbk];
        int8_t *B_char = ((int8_t*)B)+idn*k+idbk*bk;
        for(int id=0; id<bk; ++id){
          int8_t weight = A_uint1&(1<<id) ? 1 : -1;
          dot += weight*int8_t(B_char[id]);
          // if(idm == 0 && idn == 0){
          //   printf("weight %d active %d kid %d \n", weight, int8_t(B_char[id]), id);
          // }
        }
      }
      *C_int32 = dot;
      // printf("m %d n %d dot %d \n", idm, idn, dot);
    }
  }
}

struct Case{
  int m;
  int n;
  int k;
  uint32_t A_mask;
  uint8_t B_mask;
  std::string A_path;
  std::string B_path;
  std::string C_path;
};

int main(int argc, char *argv[]) {
  std::vector<Case> cases;
  cases.push_back({1, 1, 384, 0xffffffff, 0x01, "data/c0_A.bin", "data/c0_B.bin", "data/c0_C.bin"});
  cases.push_back({8, 8, 384, 0xffffffff, 0x01, "data/c1_A.bin", "data/c1_B.bin", "data/c1_C.bin"});
  // cases.push_back({231400, 1, 384, 0xff00f000, 0xff, "data/c2_A.bin", "data/c2_B.bin", "data/c2_C.bin"});
  // cases.push_back({330167, 32, 384, 0x000f00ff, 0xff, "data/c3_A.bin", "data/c3_B.bin", "data/c3_C.bin"});
  // cases.push_back({170028, 63, 384, 0x0000000f, 0xff, "data/c4_A.bin", "data/c4_B.bin", "data/c4_C.bin"});
  for(int i = 1; i < 129; ++i){
    int n = i;
    std::ostringstream A_path, B_path, C_path;
    A_path <<"data/c" << 4+i << "_A.bin";
    B_path <<"data/c" << 4+i << "_B.bin";
    C_path <<"data/c" << 4+i << "_C.bin";   
    cases.push_back({2717, n, 384, 0xffffffff, 0xff, A_path.str(), B_path.str(), C_path.str()});
  }

  for(int id = 0; id < cases.size(); ++id){
    auto c = cases[id];
    const int m = c.m, n = c.n, k = c.k;
    std::string A_path = c.A_path, B_path = c.B_path, C_path = c.C_path;
    uint32_t A_mask = c.A_mask;
    uint8_t B_mask = c.B_mask;

    uint32_t *A = new uint32_t[m*k/32];
    uint8_t *B = new uint8_t[k*n];
    uint32_t *C = new uint32_t[m*n];
    memset(A, A_mask, m*k/32);
    memset(B, B_mask, k*n);
    memset(C, 0, m*n);
    mp_gemm_w1a8_host(m, n, k, A, B, C);
    dumpbin(A, m*k/8, A_path);
    dumpbin(B, k*n, B_path);
    dumpbin(C, m*n*4, C_path);
    delete [] A;
    delete [] B;
    delete [] C;
    printf("{\"m\": %d, \"n\": %d, \"k\": %d, \"A_path\": \"%s\", \"B_path\": \"%s\", \"C_path\": \"%s\"}\n", 
      m, n, k, A_path.c_str(), B_path.c_str(), C_path.c_str());
  }
  return 0;
}